CONUS404 Spatial Aggregation

CONUS404 Spatial Aggregation#

Using Xarray, GeoPandas and Sparse

Goal: Spatially aggregate a model data variable conservatively, i.e. by exactly partitioning each grid cell into the precise region boundaries.

Approach:

Represent both the original model grid and target grid as GeoPandas GeoSeries objects (with vector geometries)
Compute their area overlay and turn it into a sparse matrix
Perform matrix multiplication on the full Xarray dataset (with a time dimension)

It is quite fast and transparent.

%xmode minimal
import xarray as xr
import geopandas as gp
import pandas as pd
import sparse

import hvplot.pandas
import hvplot.xarray
import dask
import cf_xarray

from pynhd import NLDI, WaterData
import intake
import cartopy.crs as ccrs

Exception reporting mode: Minimal

def configure_cluster(machine):
    ''' Helper function to configure cluster
    '''
    if machine == 'denali':
        from dask.distributed import LocalCluster, Client
        cluster = LocalCluster(threads_per_worker=1)
        client = Client(cluster)
    
    elif machine == 'tallgrass':
        from dask.distributed import Client
        from dask_jobqueue import SLURMCluster
        cluster = SLURMCluster(queue='cpu', cores=1, 
                               walltime="01:00:00", account="woodshole",
                               interface='ib0', memory='6GB')
        cluster.adapt(maximum_jobs=10)
        client = Client(cluster)
        
    elif machine == 'local':
        import os
        import warnings
        from dask.distributed import LocalCluster, Client
        warnings.warn("Running locally can result in costly data transfers!\n")
        n_cores = os.cpu_count() # set to match your machine
        cluster = LocalCluster(threads_per_worker=n_cores)
        client = Client(cluster)
        
    elif machine in ['esip-qhub-gateway-v0.4']:   
        import sys, os
        sys.path.append(os.path.join(os.environ['HOME'],'shared','users','lib'))
        import ebdpy as ebd
        aws_profile = 'esip-qhub'
        ebd.set_credentials(profile=aws_profile)

        aws_region = 'us-west-2'
        endpoint = f's3.{aws_region}.amazonaws.com'
        ebd.set_credentials(profile=aws_profile, region=aws_region, endpoint=endpoint)
        worker_max = 10
        client,cluster = ebd.start_dask_cluster(profile=aws_profile, worker_max=worker_max, 
                                              region=aws_region, use_existing_cluster=True,
                                              adaptive_scaling=False, wait_for_cluster=False, 
                                              worker_profile='Medium Worker', propagate_env=True)
        
    return client, cluster

Same notebook can be run on on-prem HPC or Cloud#

Here we start as Dask Cluster on the USGS HPC Denali system or on a JupyterHub deployment with Kubernetes running on Cloud

import os
if 'SLURM_CLUSTER_NAME' in os.environ:    # on prem
    dataset = 'conus404-hourly-onprem'
    machine = os.environ['SLURM_CLUSTER_NAME']
    client, cluster = configure_cluster(machine) 
else:                                     # on cloud
    dataset = 'conus404-hourly-cloud'     
    machine = 'esip-qhub-gateway-v0.4'
    client, cluster = configure_cluster(machine)

Region: us-west-2
No Cluster running.
Starting new cluster.

{}
Setting Cluster Environment Variable AWS_DEFAULT_REGION us-west-2

Setting Fixed Scaling workers=10
Reconnect client to clear cache
client.dashboard_link (for new browser tab/window or dashboard searchbar in Jupyterhub):
https://nebari.esipfed.org/gateway/clusters/dev.2e772182b32447d9bcdecc80d0faa780/status
Propagating environment variables to workers
Using environment: users/users-pangeofu

Load the feature polygons (e.g. here catchment basins)#

Load with geopandas:

%%time
# USGS gage 01482100 Delaware River at Del Mem Bridge at Wilmington De
gage_id = '01482100'
nldi = NLDI()
del_basins = nldi.get_basins(gage_id)
huc12_basins = WaterData('huc12').bygeom(del_basins.geometry[0])

CPU times: user 1.47 s, sys: 173 ms, total: 1.64 s
Wall time: 24.3 s

regions_df = huc12_basins
region_name = 'name'

Check it:

regions_df.plot(column=region_name, figsize=(10,4))

<Axes: >

../_images/3176082a918577af17604d4be4f2a8f61b24134eff2c0a1bfd008f55e1fa7067.png

All geodataframes should have a coordinate reference system. This is important (and sometimes unfamiliar to users coming from the global climate world).

crs_orig = f'EPSG:{regions_df.crs.to_epsg()}'
crs_orig

'EPSG:4326'

Key Step: Overlay the two geometries#

This is the magic of geopandas; it can calculate the overlap between the original grid and the regions. It is expensive because it has to compare 64800 grid boxes with 242 regions.

In this dataframe, the latitude and longitude values are from the grid, while all the other columns are from the regions.

overlay = grid_df.overlay(regions_df, keep_geom_type=True)

This is essentially already a sparse matrix mapping one grid space to the other. How sparse?

sparsity = len(overlay) / (len(grid_df) * len(regions_df))
sparsity

0.0024750477952660914

Let’s explore these overlays a little bit

We can verify that each basin’s area is preserved in the overlay operation.

overlay.geometry.area.groupby(overlay[region_name]).sum().nlargest(10)/1e6  # km2

name
Delaware Bay-Deep                         1073.773336
Ashokan Reservoir-Esopus Creek             205.798540
Lower Little Swatara Creek                 194.683255
Lake Wallenpaupack-Wallenpaupack Creek     180.104323
Bear Creek                                 176.743216
Bear Kill                                  174.877729
Upper Mahanoy Creek                        170.342913
Headwaters Paulins Kill River              167.072157
Black Creek                                160.267392
Snitz Creek-Quittapahilla Creek            159.798726
dtype: float64

overlay[overlay[region_name] == "Delaware Bay-Deep"].geometry.plot(edgecolor='k', aspect='equal')

<Axes: >

../_images/2a582d441b4c1538c7e150c7332b8a391d247742e8a93dcab1aaf61d6d5330f9.png

regions_df.geometry.area.groupby(regions_df[region_name]).sum().nlargest(10)

name
Delaware Bay-Deep                         1.073773e+09
Ashokan Reservoir-Esopus Creek            2.057985e+08
Lower Little Swatara Creek                1.946833e+08
Lake Wallenpaupack-Wallenpaupack Creek    1.801043e+08
Bear Creek                                1.767432e+08
Bear Kill                                 1.748777e+08
Upper Mahanoy Creek                       1.703429e+08
Headwaters Paulins Kill River             1.670722e+08
Black Creek                               1.602674e+08
Snitz Creek-Quittapahilla Creek           1.597987e+08
dtype: float64

Calculate the area fraction for each region#

This is another key step. This transform tells us how much of a country’s total area comes from each of the grid cells. This is accurate because we used an area-preserving CRS.

grid_cell_fraction = overlay.geometry.area.groupby(overlay[region_name]).transform(lambda x: x / x.sum())
grid_cell_fraction

     0.006051
     0.007286
     0.000754
     0.004720
     0.014627
          ...   
  0.001910
  0.179560
  0.180305
  0.093585
  0.000193
Length: 6196, dtype: float64

We can verify that these all sum up to one.

grid_cell_fraction.groupby(overlay[region_name]).sum()

name
Alexauken Creek-Delaware River                       1.0
Allegheny Creek-Delaware River                       1.0
Angelica Creek-Schuylkill River                      1.0
Antietam Creek                                       1.0
Appenzell Creek                                      1.0
                                                    ... 
Wrangle Brook                                        1.0
Wright Brook-Headwaters West Brach Delaware River    1.0
Wright Creek-Lehigh River                            1.0
Wyomissing Creek                                     1.0
Yaleville Brook-Susquehanna River                    1.0
Length: 451, dtype: float64

Turn this into a sparse Xarray DataArray#

The first step is making a MultiIndex

multi_index = overlay.set_index([y, x, region_name]).index
df_weights = pd.DataFrame({"weights": grid_cell_fraction.values}, index=multi_index)
df_weights

			weights
y	x	name
216039.100363	1.919902e+06	Delaware Bay-Deep	0.006051
216039.100363	1.923902e+06	Delaware Bay-Deep	0.007286
220039.100363	1.911902e+06	Delaware Bay-Deep	0.000754
	1.915902e+06	Delaware Bay-Deep	0.004720
	1.919902e+06	Delaware Bay-Deep	0.014627
...	...	...	...
616039.100363	1.859902e+06	Mine Kill	0.001910
	1.863902e+06	Mine Kill	0.179560
	1.867902e+06	Mine Kill	0.180305
	1.871902e+06	Mine Kill	0.093585
	1.875902e+06	Mine Kill	0.000193

6196 rows × 1 columns

We can bring this directly into xarray as a 1D Dataset.

import xarray as xr
ds_weights = xr.Dataset(df_weights)

Now we unstack it into a sparse array.

weights_sparse = ds_weights.unstack(sparse=True, fill_value=0.).weights

Here we can clearly see that this is a sparse matrix, mapping the input space (lat, lon) to the output space (SOVEREIGNT).

CONUS404 Spatial Aggregation

Contents

CONUS404 Spatial Aggregation#

Same notebook can be run on on-prem HPC or Cloud#

Load the feature polygons (e.g. here catchment basins)#

Open the gridded model dataset (e.g. here CONUS404)#

Key Step: Overlay the two geometries#

Calculate the area fraction for each region#

Turn this into a sparse Xarray DataArray#

Perform Matrix Multiplication#

Variables at the same location on the grid can use the same weights#

Explore timeseries data by region:#

Explore the mean var by region#

Be a good citizen and shut down the cluster#